import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import math
import matplotlib.pyplot as plt

noise_dim = 200


class Generator(keras.Model):
    def __init__(self):
        super(Generator, self).__init__()

        self.dense = layers.Dense(4 * 4 * 2048, input_shape=(noise_dim,),
                                  use_bias=False,
                                  activation="relu",
                                  kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.reshape = layers.Reshape((4, 4, 2048))

        self.conv1 = layers.Conv2DTranspose(filters=1024,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.bn1 = layers.BatchNormalization()
        self.relu1 = layers.ReLU()

        self.conv2 = layers.Conv2DTranspose(filters=512,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.bn2 = layers.BatchNormalization()
        self.relu2 = layers.ReLU()

        self.conv3 = layers.Conv2DTranspose(filters=256,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.bn3 = layers.BatchNormalization()
        self.relu3 = layers.ReLU()

        self.conv4 = layers.Conv2DTranspose(filters=128,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.bn4 = layers.BatchNormalization()
        self.relu4 = layers.ReLU()

        self.conv5 = layers.Conv2DTranspose(filters=64,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02))
        self.bn5 = layers.BatchNormalization()
        self.relu5 = layers.ReLU()

        self.conv6 = layers.Conv2DTranspose(filters=3,
                                            kernel_size=5,
                                            strides=2,
                                            padding="same",
                                            kernel_initializer=keras.initializers.random_normal(mean=0.0, stddev=0.02),
                                            activation="tanh")

    def call(self, input):
        d1 = self.dense(input)
        r1 = self.reshape(d1)

        x_conv1 = self.conv1(r1)
        b1 = self.bn1(x_conv1)
        re1 = self.relu1(b1)

        x_conv2 = self.conv2(re1)
        b2 = self.bn2(x_conv2)
        re2 = self.relu2(b2)

        x_conv3 = self.conv3(re2)
        b3 = self.bn3(x_conv3)
        re3 = self.relu3(b3)

        x_conv4 = self.conv4(re3)
        b4 = self.bn4(x_conv4)
        re4 = self.relu4(b4)

        x_conv5 = self.conv5(re4)
        b5 = self.bn5(x_conv5)
        re5 = self.relu5(b5)

        x_conv6 = self.conv6(re5)
        return x_conv6


G = Generator()
G.build(input_shape=(1, noise_dim))
G.load_weights("save/myG.h5")

for i in range(40):
    noise = tf.random.normal([1, noise_dim])
    image = G(noise, training=False)
    one = tf.ones([256, 256, 3])
    plt.margins(0, 0)
    plt.imshow((image[0, :, :, :] + one) / 2.0)
    plt.axis('off')
    plt.savefig("result/" + "test" + str(i) + ".png", bbox_inches='tight', pad_inches=0.0, dpi=69.5)  # 256x256像素
    plt.close()
